import re
from langchain.prompts import (
    SystemMessagePromptTemplate,
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain_experimental.plan_and_execute import (
    PlanAndExecute,
    load_agent_executor,
    load_chat_planner,
)


from pydantic import BaseModel, Field
from langchain_core.output_parsers import BaseOutputParser
import os
from langchain.agents import tool
import json

os.environ["LANGSMITH_TRACING"] = "true"
# os.environ["LANGSMITH_API_KEY"] = getpass.getpass()
os.environ["LANGSMITH_API_KEY"] = "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"

# from langchain_openai import ChatOpenAI

# llm = ChatOpenAI(
#     # openai_api_key="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",  # gtihub
#     openai_api_key="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",  # 谷歌
#     openai_api_base="https://openrouter.ai/api/v1",
#     # model_name="deepseek/deepseek-chat-v3-0324:free",
#     model_name="qwen/qwq-32b:free",
# )

from langchain_ollama import ChatOllama

llm = ChatOllama(
    base_url="http://222.210.44.235:11434",
    model="qwen2.5:32b",
    temperature=0,
    num_ctx=32768,
)


plan_tmp = """\
    首先请你理解这个问题，然后设计解决这个问题的计划，设计计划过程中必须遵守以下规则：
    1、计划的表头必须使用“plan:”开始，然后是步骤的编号列表
    2、以最少的计划步骤准确的完成计划设计
    3、设计计划时优先考虑使用工具列表中的工具完成计划任务，工具列表：{Tools}
    4、设计计划时需考虑计划的依赖性和连贯性，尽量使用原始问题中的相关信息调用对应工具完成计划任务
    5、调用工具时严格遵守工具的参数要求，具体信息参考工具列表详细描述
    6、计划步骤中严禁出现重复计划或者相似的重复计划
    7、计划的最后一步需考虑之前的计划是否能解决原始问题
    8、在计划的末尾必须使用“<END_OF_PLAN>”表示计划结束
    9、相关工具返回结果中“<END_RUN>”表示结束当前全部计划
    
    以下为提供设计计划的工具列表详细描述：
    {Tools_Describe}
"""


def toData(desc):
    sys_tmp = """\
        角色设定：参数分析助手
        职责：首先分析输入数据，然后根据输出数据格式提取对应参数，最后根据输出数据格式返回结果。

        输出数据：
        ```json
        [{{"label": "名称", "value": "数值"}}]
        ```
    """

    human_tmp = """\
        输入内容：
        {desc}
    """

    sys_msg = SystemMessagePromptTemplate.from_template(sys_tmp)
    human_msg = HumanMessagePromptTemplate.from_template(human_tmp)
    chat_msg = ChatPromptTemplate.from_messages([sys_msg, human_msg])

    data = ""
    try:
        chain = (chat_msg | llm | jsonOutputParser()).with_retry()
        data = chain.invoke(
            {
                "desc": desc,
            },
        )
        print("data", data)
    except Exception as e:
        print(e)
        return "数据分析失败，请重试"

    return data


@tool()
def overview(*args, **kwargs):
    """
    概览组件，突出显示关键指标的核心摘要
    :return 组件JSON格式
    """

    res = {
        "title": "<概览标题>",
        "type": "概览",
        "data": [{"label": "<标签>", "value": "<数值>"}],
    }

    return json.dumps(res, ensure_ascii=False)


class jsonOutputParser(BaseOutputParser[str]):
    def parse(self, text: str) -> str:
        print("htmlOutputParser:", text)
        pattern = r"```json(.*?)```"
        matches = re.findall(pattern, text, re.DOTALL)
        if matches:
            return matches[0].strip()
        else:
            return "未找到HTML代码"

    @property
    def _type(self) -> str:
        return ""


GenerateSocHtmlDesc = """\
提取适合生成柱状图的文本段落

要求：
1. 文本必须包含至少2个可对比的类别及其对应数值
2. 数值需为明确数字（如100、25%、$300等）
"""


class barChartInput(BaseModel):
    title: str = Field(..., description="生成柱状图组件的标题")
    desc: str = Field(
        ...,
        description="适合柱状图的描述段落（如：漏洞等级高危、中危、低危数量）",
    )


@tool(description=GenerateSocHtmlDesc, args_schema=barChartInput)
def barChart(*args, **kwargs) -> str:
    title = ""
    if "title" in kwargs:
        title = kwargs["title"]
    else:
        return "输入参数错误，根据工具参数要求输入title参数"

    desc = ""
    if "desc" in kwargs:
        desc = kwargs["desc"]
    else:
        return "输入参数错误，根据工具参数要求输入desc参数"

    data = toData(desc)

    res = {
        "title": title,
        "type": "柱状图",
        "data": data,
    }

    return json.dumps(res, ensure_ascii=False)


@tool()
def piebart(*args, **kwargs) -> str:
    """
    柱状图组件 - 展示整体中各部分的比例关系
    :return 组件JSON格式
    """

    res = {
        "title": "<组件标题>",
        "type": "饼图",
        "data": [{"label": "<标签>", "value": "<数值>"}],
    }

    return json.dumps(res, ensure_ascii=False)


@tool()
def buidlJSon(*args, **kwargs) -> str:
    """
    将多个可视化组件的数据智能组合成标准化的仪表板数据结构。

    要求：
    返回JSON格式必须保持一致

    返回示例：
    {{
        "header": {{
            "title": "网络安全态势分析仪表板",
            "description": "资产事件核心指标与攻击类型分布可视化"
        }},
        "components": [
            {{
                "type": "piebart",
                "data": [
                    {{
                        "label": "恶意软件",
                        "value": 45
                    }}
                ],
                "title": "攻击类型分布比例"
            }}
        ]
    }}

    :return 合并后的JSON格式参考
    """

    data = {
        "header": {
            "title": "<根据内容生成标题>",
            "description": "<根据内容生成描述>",
        },
        "components": [
            {
                "title": "<组件标题>",
                "type": "<组件类型>",
                "data": "<组件数据>",
            }
        ],
    }

    return json.dumps(data, ensure_ascii=False)


web_json = []


@tool()
def saveJson(
    json_data: str,
) -> str:
    """
    将网页JSON数据保存到本地

    Args:
        json_data: 要保存的JSON数据字符串

    Returns:
        str: 保存结果消息
    """
    # 将新数据添加到列表
    web_json.append(json_data)

    # 将所有数据连接起来写入文件（根据你的需求可能需要调整）
    with open("web.json", "w", encoding="utf-8") as file:
        file.write("\n".join(web_json))  # 假设每条数据用换行符分隔

    return "保存网页成功"


# 工具列表
# tools = [overview, piebart, barChart, buidlJSon, saveJson]
tools = [barChart, barChart, saveJson]

ret = barChart.invoke({"title": "test", "desc": "aaaa"})
print(ret)
exit()
# tools = [barChart, saveJson]
# tools = [test]
tool_strings = []
for t in tools:
    args_schema = re.sub("}", "}}", re.sub("{", "{{", str(t.args)))
    tool_strings.append(f"{t.name}: {t.description}, args: {args_schema}")
    formatted_tools = "\n".join(tool_strings)
    tool_names = ", ".join([t.name for t in tools])

sys_prompt_tmp = SystemMessagePromptTemplate.from_template(plan_tmp)
chat_sys_tmp = ChatPromptTemplate.from_messages([sys_prompt_tmp])
sys_prompt = chat_sys_tmp.format(Tools=tool_names, Tools_Describe=formatted_tools)

planner = load_chat_planner(llm, system_prompt=sys_prompt)
executor = load_agent_executor(llm, tools, verbose=True)


summary = """
### 网络安全摘要信息

#### 资产网络攻击统计相关数据

- **资产 ID**: `0fd4b9d96d6f5ecab6c6a8b6809f5ae3`
- **资产名称**: `taomingwei-node3-ubuntu`
- **总事件数量**: 13 个
- **告警时间范围**：
  - 第一条告警时间：2022 年 4 月 3 日 09:52:39
  - 最新告警时间：2025 年 4 月 8 日 09:12:25

#### 资产网络攻击类型分布

- **攻击 ID 及描述**：

  - T1564.001：通过命令行创建隐藏文件和目录，告警级别为 3。
  - T1078.003：通过 Grep 命令查找系统上包含 Password 关键字的文件，告警级别为 3。
  - T1082：可疑 Modprobe 文件事件，告警级别为 4。
  - T1059：
    - 检测到 Linux tmp 目录程序运行和检测到 Linux 可疑 shell 命令，告警级别分别为 4（第 3、5 条数据）。
    - Linux 受限 Shell 通过 Linux 二进制突破，告警级别为 3。
    - 检测到 Linux perl 反弹 shell2，威胁级别: 4
  - T1027.001 (填充二进制文件)
  - T1070.002 (清除 Linux 系统日志)

- **攻击阶段**：
  - `TA0007`（第 2 条数据）
  - `TA0002`（第 3、4、5 条数据，以及检测到 Linux perl 反弹 shell 事件）
  - `TA0004`（填充二进制文件）
  - `TA0005`（清除 Linux 系统日志）

#### 资产漏洞相关情况

- **检测方式**：
  - 隐藏文件目录检测：通过命令行创建隐藏文件和目录。
  - 敏感命令执行检测：使用 Grep 命令查找系统上包含 Password 关键字的文件。
  - 发现系统信息检测：检测 modprobe 配置文件中涉及内核模块的文件事件。
  - 检测到 Linux tmp 目录程序运行（第 3 条数据）
  - Linux 受限 Shell 通过 Linux 二进制突破（第 4 条数据）
  - 检测到 Linux 可疑 shell 命令（第 5 条数据）
  - 填充二进制文件: 检测通过 DD 指令填充二进制的行为。
  - 清除 Linux 系统日志: 黑客通常在入侵结束后对入侵痕迹进行清理，例如对所有可能记录其攻击行为的日志进行删除操作。

#### 资产风险相关情况

- **告警级别**：
  - 第 3、5 条数据：告警级别为 4。
  - 第 4 条数据：告警级别为 3。
  - 填充二进制文件: 级别 3 (第 6 条)
  - 清除 Linux 系统日志: 级别 3 (第 7 条)
  - 检测到 Linux 可疑命令: 级别 4 (第 8 条)

#### 资产安全防护状态

- **修复建议**：
  - 排查和结束相关可疑进程。
  - 排查和删除相关文件。
  - 排查和删除相关系统设置。
  - 排查和删除相关注册表项。
  - 填充二进制文件:
    1. 通过任务管理排查和结束相关可疑进程。
    2. 通过资源管理器排查和删除相关文件。
    3. 排查和删除相关可疑系统设置。
    4. 排查和删除相关可疑注册表项。
  - 清除 Linux 系统日志:
    1. 确认删除的日志文件是否可疑，排除正常的运维与文件更新操作。
    2. 确认操作可疑的情况下，进一步排查导致文件删除的进程名以及对应的用户名，必要时结束进程。
    3. 对可疑进程的后续操作进行排查。
  - 检测到 Linux 可疑命令:
    排查命令上下文分析命令是否恶意，若为恶意及时结束相关进程。

### 总结

资产 `taomingwei-node3-ubuntu` 遭遇了多次网络安全事件：

1. 检测到 Linux 可疑命令。
2. 检测到 Linux perl 反弹 shell。

建议立即排查相关网络连接并采取措施防止进一步的恶意活动。

"""


try:
    agent = PlanAndExecute(planner=planner, executor=executor, verbose=True)
    input_old = """\
        任务要求：
        1、深入分析输入内容，为输入内容生成对应的可视化组件
        2、将多个可视化组件的数据智能组合成标准化的仪表板数据结构，严格按照工具的JONS格式组合
        3、智能组合JSON数据，生成可视化仪表板
        4、JSON生成完以后保存JSON，保存JSON的格式需要维持组件的格式，不要自行修改

        输入内容：
        {summary}

    """
    # input = """\
    #     任务要求：
    #     1、深入分析输入内容，将输入内容转换为锻炼，送入工具
    #     2、将多个可视化组件的数据智能组合成标准化的仪表板数据结构，严格按照工具的JONS格式组合
    #     3、智能组合JSON数据，生成可视化仪表板
    #     4、JSON生成完以后保存JSON，保存JSON的格式需要维持组件的格式，不要自行修改

    #     输入内容：
    #     {summary}
    # """

    input = """\
        任务要求：
        1、深入分析输入内容，将适合段落送入工具提取（可以提取多次，保证数据的完整性）

        输入内容：
        {summary}
    """

    input = input.format(summary=summary)

    res = agent.invoke(input, verbose=True)

    # print("输出结果", res.content)
except Exception as e:
    print("异常", e, input)
